import platgo as pg
import numpy as np


class ArmSearch:

    def __init__(self, problem: pg.Problem, beta, sigma):
        self.problem = problem
        self.beta = beta
        self.sigma = sigma

    def get_SL(self, pop: pg.Population, dk: np.ndarray, gk: np.ndarray) -> np.ndarray:
        """
        通过Armijo搜索求个体的步长
        :param pop: 种群
        :param dk: 搜索方向
        :param gk: 函数在各个个体对应的梯度
        :return: 每个个体的步长
        """

        mk = 0
        for i in range(20):
            pops = pg.Population(decs=pop.decs + (self.beta ** i) * dk)
            self.problem.cal_obj(pops)
            if pops.objv <= (pop.objv + self.sigma * self.beta ** i * np.dot(gk, dk.T)):
                mk = i
                break
        return mk

    def __call__(self, pop: pg.Population, dk: np.ndarray, gk: np.ndarray):
        return self.get_SL(pop, dk=dk, gk=gk)



